(*
  HDFS Digital Logic Hardware Design Utility Library (hdfslib.dll)
  Copyright (C) 2006 Andy Ray.

  This library is free software; you can redistribute it and/or
  modify it under the terms of the GNU Lesser General Public
  License as published by the Free Software Foundation; either
  version 2.1 of the License, or (at your option) any later version.

  This library is distributed in the hope that it will be useful,
  but WITHOUT ANY WARRANTY; without even the implied warranty of
  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
  Lesser General Public License for more details.

  You should have received a copy of the GNU Lesser General Public
  License along with this library; if not, write to the Free Software
  Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301  USA
*)

(** General utilities useful for hardware design. *)
module DigitalLogic.Util

open DigitalLogic
open Numeric.Ops
open Circuit
open Signal
open List
open Big_int

exception Reduce_error of string

(*****************************************************************************)
(*****************************************************************************)

(** given a (positive) value calculates the required bit width ie ceil(log2(x)) *)
let rec clog2 x = match x with 0 | 1 -> 1 | x -> 1 + (clog2 (x/2)) 

(** 2 ^ n in integer arithmetic *)
let rec pow2 n = 2 * (if n<=1 then 1 else pow2 (n-1))

(** creates the list [ 0; 1; 2; .. ; (n-1) ] (note that the new list comprehension syntax in F# can perform this task) *)
let range n = 
    let rec r0 n i = if n = i then [] else i :: (r0 n (i+1)) in
    r0 n 0 

(** select a range from the list *)
let lselect l lo hi = 
    let rec ls l idx lo hi = 
        if idx > hi then []
        else if idx < lo then ls (tl l) (idx+1) lo hi
        else (hd l) :: (ls (tl l) (idx+1) lo hi) in 
    ls l 0 lo hi 

(** Selects the even elements from a list *)
let leven l = 
  let rec r l n = 
    match l with
    | [] -> []
    | hd :: tl ->
      if (n &&& 1) = 0 
      then hd :: (r tl (n+1))
      else (r tl (n+1)) in
  r l 0
      
(** Selects the odd elements from a list *)
let lodd l = 
  let rec r l n = 
    match l with
    | [] -> []
    | hd :: tl ->
      if (n &&& 1) = 1 
      then hd :: (r tl (n+1))
      else (r tl (n+1)) in
  r l 0


(*****************************************************************************)
(* vector operations *)
(*****************************************************************************)

(** resize vec depending on sign *)
let resize sign vec newwid = 
    let oldwid = width vec in
    let extend vec = (repeat (if sign = Signed then msb vec else gnd) (newwid-oldwid)) ++ vec in
    let shrink vec = select vec (newwid-1) 0 in
    if oldwid = newwid then vec
    else if oldwid < newwid then extend vec
    else shrink vec

(** resize vec adding 0's at the msb as required *)
let uresize = resize Signed

(** resize vec repeating the msb to preserve sign *)
let sresize = resize Unsigned

(** number of bits required to represent the given value, which may be signed *)
let num_bits v = 
    if v < 0 then (clog2 (abs (v+1))) + 1
    else clog2 v

(** absolute value *)
let abs a = mux2 (msb a) (- a) a 

(** pad vector (at right handside ie below lsb) up to a power of two with the given constant (which should be 1 bit) *)
let pad_right_pow2 v c =
  let len = width v in
  let log = clog2 len in
  let llen = pow2 log in
  if llen - len > 0 then
    v ++ (repeat c (llen-len))
  else v

(* should these be moved to design.ml? *)

(** unsigned addition.  Output is 1 bit larger than parameters so cannot overflow *)
let (+:+) a b = (gnd ++ a) + ((gnd ++ b) : Signal)

(** signed addition.  Output is 1 bit larger than parameters so cannot overflow *)
let (+:-) a b = ((msb a) ++ a) + (((msb b) ++ b) : Signal)     

(** unsigned subtaction.  Output is 1 bit larger than parameters so cannot overflow *)
let (-:+) a b = (gnd ++ a) - ((gnd ++ b) : Signal)                

(** signed subtaction.  Output is 1 bit larger than parameters so cannot overflow *)
let (-:-) a b = ((msb a) ++ a) - (((msb b) ++ b) : Signal)

(** Counts from 0 to (max-1) then from zero again.  If max == 1<<n, then the comparator is not generated and overflow arithmetic used instead *)
let mod_counter max c = 
  let w = width c in
  let lmax = 1 <<< w in
  if lmax = (max + 1) 
  then c + 1
  else mux2 (c ==: max) (zero w) (c + 1)

(** creates a tree of operations.  The arity of the operator is configurable *)
let rec tree arity ops l =
  let split l n =
     let (lh,ll,_) = List.fold_left (fun (l0,l1,m) e ->
       if m < n then ((e::l0),l1,m+1) else (l0,e::l1,m+1)) ([],[],0) l in
     (List.rev lh, List.rev ll) in
  let rec t0 l =
    let l0,l1 = split l arity in
    if l1 = [] then [ ops l0 ]
    else (ops l0) :: (t0 l1) in
  match l with
  | [] -> failwith "Invalid list given to tree"
  | [a] -> a
  | _ -> tree arity ops (t0 l) 

(** creates a binary tree of operations *)
let binary_tree uop bop = tree 2 
  (function [a] -> uop a | [a;b] -> bop a b | _ -> failwith "Binary tree expects no more than two arguments per operation")

(*****************************************************************************)
(*****************************************************************************)

(** Builds a pipeline of registers. *)
let rec pipeline n enable d = 
  if n=0 then d
  else regc enable (pipeline (n-1) enable d)

(** Same as pipeline except that reset is not specified.  This allows it to be inferred as SRL16 in Virtex FPGAs. *)
let rec pipeline_srl n enable d = 
  if n=0 then d
  else reg clock empty empty enable (pipeline_srl (n-1) enable d)

(*****************************************************************************)
(*****************************************************************************)

(** binary encoding to one hot encoding.  N cases where N is <= (2 ** width s).  Unused cases driven to 0 *)
let binary_to_one_hot addr n = 
  if n <= 0 then empty
  else (
    assert (n <= (1 <<< (width addr)));
    let wid = width addr in
    let one_hot = b_wire0 n in
    behave [
      b_switch (addr) (
        List.map 
          (fun i -> b_case (consti wid i) [ one_hot $== const1h n i; ])
            [ 0 .. (n-1) ]
      )
    ];
    one_hot.q
  )

(*****************************************************************************)
(*****************************************************************************)

(* Build roms based on the given function operating on 1 or 2 arguments *)

let op_rom' fn in_bits out_bits = 
  let rec aux i = 
    if i = (1 <<< in_bits) - 1 then []
    else consti out_bits (fn i) :: aux (i+1)
  in
  aux 0

let op_rom_1 fn in_bits out_bits = 
  fun addr -> mux addr (op_rom' fn in_bits out_bits)

let op_rom_2 fn in_bits_0 in_bits_1 out_bits = 
  let rec aux i = 
    if i = (1 <<< in_bits_0) - 1 then []
    else op_rom' (fn i) in_bits_1 out_bits @ aux (i+1)
  in
  fun (addr0:Signal) (addr1:Signal) -> mux (addr0 ++ addr1) (aux 0)

let rom_add (a:Signal) (b:Signal) = op_rom_2 (+) a.width b.width (max a.width b.width)
let rom_sub (a:Signal) (b:Signal) = op_rom_2 (-) a.width b.width (max a.width b.width)
let rom_mul (a:Signal) (b:Signal) = op_rom_2 ( * ) a.width b.width (a.width + b.width)

(*****************************************************************************)
(*****************************************************************************)

(** Binary reflected gray code *)
let make_gray_code_strings l = 
  let gray1 = [ "0"; "1" ] in
  if l <= 0 then failwith "Gray code with less than 1 bit not possible"
  else
    let rec build_gray_code n codes = 
      if n=l then codes
      else
        build_gray_code (n+1) 
          ((map ((^) "0") codes) @ (map ((^) "1") (rev codes))) in
    build_gray_code 1 gray1

(** Constructs a statemachine, where each state has a string name.  The state signal, and a function to look up a state value given it's name is returned. *)
let make_states clock reset enable states = 
  let num_states = List.length states in
  let log_states = clog2 (num_states-1) in
  let states = mapi (fun i x -> (i,x)) states in
  let map = List.fold_left (fun map (idx,str) -> Map.add str (consti log_states idx) map) Map.empty states in
  b_reg clock reset empty enable log_states, (fun str -> Map.find str map)

(** Constructs a statemachine using a one-hot encoding *)
let make_states_onehot clock reset enable states = 
  let num_states = List.length states in
  let states = mapi (fun i x -> (i,x)) states in
  let map = List.fold_left (fun map (idx,str) -> Map.add str (const1h num_states idx) map) Map.empty states in
  b_reg clock reset (const1h num_states 0) enable num_states, (fun str -> Map.find str map)


(** Constructs a statemachine using a gray encoding *)
let make_states_gray clock reset enable states = 
  let num_states = List.length states in
  let log_states = clog2 (num_states-1) in
  let gray_strings = lselect (make_gray_code_strings log_states) 0 (num_states-1) in
  let states = map2 (fun x g -> (x,g)) states gray_strings in
  let map = List.fold_left (fun map (str,gray) -> Map.add str (constb gray) map) Map.empty states in
  b_reg clock reset empty enable log_states, (fun str -> Map.find str map)

(*****************************************************************************)
(*****************************************************************************)

type StateEncoding = SECount | SEOneHot | SEGray

(* Encoding of state machine enumerations, up to 32 bits, for use with behavioural blocks *)
(* Note: might be sensible to allow multiple state variables be generated - this makes it more like a type *)
type StateEnum = 
  {
    states : (string * Signal) list;
    state_var : BehaveAssignTarget;
  }
  with
  
    static member make clock reset enable encoding states = 
      let num_states = List.length states in
      let log_num_states = ubits (num_states-1) in
      let codes, state_width = 
        match encoding with
        | SECount -> [ for i in { 0 .. num_states - 1 } -> consti log_num_states i ], log_num_states
        | SEOneHot -> [ for i in { 0 .. num_states - 1 } -> const1h num_states i ], num_states
        | SEGray -> map constb (lselect (make_gray_code_strings log_num_states) 0 (num_states-1)), log_num_states
      in
      {
        states = List.combine states codes;
        state_var = b_reg clock reset (List.hd codes) enable state_width;
      }
    
    member x.signal_of_state str = 
      try List.assoc str x.states 
      with _ -> failwith ("Couldn't look up state " ^ str)

    member x.Item with get(state:string) = x.signal_of_state state
    
    (* given the state value encoded as a binary string, return the state name *)
    member x.state_of_binary_string str = 
      let rec find s = 
        match s with
        | [] -> failwith ("Couldnt look up state with value " ^ str)
        | (name,signal)::tl ->
          if str = string_of_const signal then name
          else find tl
      in
      find x.states
      
    member x.v = x.state_var
    member x.q = x.v.q
    
    static member (==:) ((state:StateEnum), (s:string)) = (==:) state.q (state.[s] : Signal)
    static member (/=:) ((state:StateEnum), (s:string)) = (/=:) state.q (state.[s] : Signal)
    static member ($==) ((state:StateEnum), (s:string)) = ($==) state.v (state.[s] : Signal)

  end
